from deep_translator import GoogleTranslator
import os
from pathlib import Path
import shutil
import pypdf
import re
from io import BytesIO
import zipfile
from bs4 import BeautifulSoup

def translate_pdf_file(input_path, output_path=None, source_lang='auto', target_lang='zh-CN', from_page=1, to_page=None):
    """
    翻译PDF文件（支持多页，支持页码范围选择，保持PDF格式）
    
    参数:
        input_path (str): 输入PDF文件路径
        output_path (str, optional): 输出文件路径，默认为None（自动生成）
        source_lang (str): 源语言代码，默认为auto自动检测
        target_lang (str): 目标语言代码，默认为zh-CN
        from_page (int): 起始页码（从1开始），默认为1
        to_page (int, optional): 结束页码（包含），默认为None（表示最后一页）
        
    返回:
        tuple: (翻译后的文本, 输出文件路径, 文本文件路径)
    """
    try:
        # 确保输入文件存在
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"输入文件不存在: {input_path}")
        
        # 检查页码范围
        if from_page < 1:
            raise ValueError(f"起始页码必须大于等于1，当前值: {from_page}")
        
        # 如果未指定输出路径，则自动生成
        if output_path is None:
            input_file = Path(input_path)
            filename = input_file.stem
            page_range = ""
            if from_page > 1 or to_page is not None:
                page_range = f"_p{from_page}"
                if to_page is not None:
                    page_range += f"-{to_page}"
            output_path = os.path.join("output_files", f"{filename}{page_range}_{target_lang}.pdf")
        elif not output_path.lower().endswith('.pdf'):
            output_path += '.pdf'
        
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 创建翻译器
        translator = GoogleTranslator(source=source_lang, target=target_lang)
        
        # 准备文本输出路径以保存所有翻译文本
        txt_output_path = output_path.replace('.pdf', '.txt')
        all_translated_text = ""
        
        # 打开PDF文件
        with open(input_path, 'rb') as pdf_file:
            pdf_reader = pypdf.PdfReader(pdf_file)
            total_pages = len(pdf_reader.pages)
            
            # 处理页码范围
            if to_page is None or to_page > total_pages:
                to_page = total_pages
            
            if from_page > total_pages:
                raise ValueError(f"起始页码超出PDF总页数。PDF共有 {total_pages} 页，指定的起始页码为 {from_page}")
            
            print(f"正在处理PDF文件，共 {total_pages} 页，翻译范围: {from_page} - {to_page} 页...")
            
            # 简化处理方式：首先复制原始PDF到输出路径
            shutil.copy2(input_path, output_path)
            
            # 逐页提取和翻译文本（仅用于生成文本文件和返回翻译文本）
            for page_num in range(from_page - 1, to_page):
                # 获取当前页
                page = pdf_reader.pages[page_num]
                # 提取文本
                page_text = page.extract_text()
                
                if page_text.strip():  # 确保页面有内容
                    try:
                        # 分段翻译，处理长文本
                        translated_page = translate_long_text(page_text, translator)
                        all_translated_text += f"\n\n--- 第 {page_num + 1} 页 ---\n\n"
                        all_translated_text += translated_page
                        print(f"已完成第 {page_num + 1} 页翻译")
                    except Exception as e:
                        all_translated_text += f"\n\n--- 第 {page_num + 1} 页（翻译失败）---\n\n"
                        all_translated_text += f"翻译错误: {str(e)}\n原文: {page_text[:100]}..."
                        print(f"第 {page_num + 1} 页翻译失败: {str(e)}")
        
        # 将翻译结果保存到文本文件
        with open(txt_output_path, 'w', encoding='utf-8') as f:
            f.write(all_translated_text)
        
        # print(f"PDF文件翻译完成，原始PDF文件已保存到: {output_path}")
        print(f"翻译文本已保存到: {txt_output_path}")
        return all_translated_text, output_path, txt_output_path  # 返回文本文件路径
    
    except Exception as e:
        raise Exception(f"PDF文件翻译失败: {str(e)}")

def translate_long_text(text, translator, max_chunk_size=4000, retry_attempts=3):
    """
    分段翻译长文本
    
    参数:
        text (str): 要翻译的文本
        translator: GoogleTranslator实例
        max_chunk_size (int): 每段最大字符数
        retry_attempts (int): 失败时的重试次数
    
    返回:
        str: 翻译后的完整文本
    """
    # 确保文本不为空
    if not text or not text.strip():
        return ""
        
    # 移除过多的空白字符，以减少翻译负载
    text = re.sub(r'\s+', ' ', text).strip()
    
    # 如果文本较短，直接翻译
    if len(text) <= max_chunk_size:
        for attempt in range(retry_attempts):
            try:
                return translator.translate(text)
            except Exception as e:
                if attempt < retry_attempts - 1:
                    import time
                    wait_time = (attempt + 1) * 2  # 逐步增加等待时间
                    print(f"翻译失败，正在重试 ({attempt+1}/{retry_attempts})，等待 {wait_time} 秒...")
                    time.sleep(wait_time)
                    continue
                raise e  # 如果所有尝试都失败，则抛出异常
    
    # 分段处理长文本
    chunks = []
    current_chunk = ""
    
    # 试图按较好的分隔符拆分文本（段落、句子）
    paragraph_split = re.split(r'(\n{2,}|\r\n{2,})', text)  # 分段
    
    for para in paragraph_split:
        # 如果当前段落加上当前块不会超过限制，直接添加
        if len(current_chunk) + len(para) <= max_chunk_size:
            current_chunk += para
            continue
            
        # 否则，需要拆分当前段落
        if current_chunk:
            chunks.append(current_chunk.strip())
            current_chunk = ""
            
        # 如果段落本身就很长，需要进一步拆分
        if len(para) > max_chunk_size:
            # 按句子分割
            sentences = re.split(r'(?<=[.!?。！？])\s+', para)
            
            for sentence in sentences:
                # 如果当前句子加上当前块不会超过限制
                if len(current_chunk) + len(sentence) <= max_chunk_size:
                    current_chunk += sentence + " "
                else:
                    # 如果当前块不为空，添加到chunks
                    if current_chunk:
                        chunks.append(current_chunk.strip())
                    
                    # 如果句子本身就很长，需要强制分割
                    if len(sentence) > max_chunk_size:
                        # 按照最大长度强制分割
                        for i in range(0, len(sentence), max_chunk_size):
                            chunks.append(sentence[i:i+max_chunk_size])
                    else:
                        current_chunk = sentence + " "
        else:
            current_chunk = para
    
    # 添加最后一个块
    if current_chunk:
        chunks.append(current_chunk.strip())
    
    # 翻译每个块并合并
    translated_chunks = []
    for i, chunk in enumerate(chunks):
        if i > 0 and i % 3 == 0:  # 每翻译3个块添加延迟
            import time
            time.sleep(1)
            
        for attempt in range(retry_attempts):
            try:
                translated_chunk = translator.translate(chunk)
                translated_chunks.append(translated_chunk)
                break  # 翻译成功，跳出重试循环
            except Exception as e:
                if attempt < retry_attempts - 1:
                    import time
                    wait_time = (attempt + 1) * 2  # 逐步增加等待时间
                    print(f"块 {i+1}/{len(chunks)} 翻译失败，正在重试 ({attempt+1}/{retry_attempts})，等待 {wait_time} 秒...")
                    time.sleep(wait_time)
                    continue
                # 所有重试失败，添加错误信息并继续下一块
                print(f"警告: 块 {i+1}/{len(chunks)} 翻译失败: {str(e)}")
                translated_chunks.append(f"[翻译失败] 原文: {chunk[:100]}...")
                break
    
    # 合并所有翻译后的块
    return " ".join(translated_chunks)

def translate_docx_file(input_path, output_path=None, source_lang='auto', target_lang='zh-CN'):
    """
    翻译DOCX文件（保持DOCX格式）
    
    参数:
        input_path (str): 输入DOCX文件路径
        output_path (str, optional): 输出文件路径，默认为None（自动生成）
        source_lang (str): 源语言代码，默认为auto自动检测
        target_lang (str): 目标语言代码，默认为zh-CN
        
    返回:
        tuple: (翻译后的文本, 输出文件路径, 文本文件路径)
    """
    try:
        # 确保输入文件存在
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"输入文件不存在: {input_path}")
        
        # 如果未指定输出路径，则自动生成
        if output_path is None:
            input_file = Path(input_path)
            filename = input_file.stem
            output_path = os.path.join("output_files", f"{filename}_{target_lang}.docx")
        elif not output_path.lower().endswith('.docx'):
            output_path += '.docx'
        
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 先复制原始文件到输出路径，以防万一文本处理失败
        shutil.copy2(input_path, output_path)
        
        # 创建翻译器
        translator = GoogleTranslator(source=source_lang, target=target_lang)
        
        # 准备文本输出路径
        txt_output_path = output_path.replace('.docx', '.txt')
        all_translated_text = ""
        
        # 使用python-docx库处理文档
        try:
            import docx
            
            # 打开复制后的文档
            doc = docx.Document(output_path)
            
            # 翻译段落
            for para in doc.paragraphs:
                if para.text.strip():
                    original_text = para.text
                    try:
                        translated_text = translate_long_text(original_text, translator)
                        all_translated_text += translated_text + "\n\n"
                    except Exception as e:
                        print(f"警告: 段落翻译失败: {str(e)[:100]}...")
                        # 保留原始文本
                        all_translated_text += f"[翻译失败] {original_text}\n\n"
            
            # 翻译表格内容（仅文本记录，不修改原文档）
            for table in doc.tables:
                for row in table.rows:
                    for cell in row.cells:
                        if cell.text.strip():
                            original_text = cell.text
                            try:
                                translated_text = translate_long_text(original_text, translator)
                                all_translated_text += translated_text + " | "
                            except Exception as e:
                                print(f"警告: 表格单元格翻译失败: {str(e)[:100]}...")
                                # 保留原始文本
                                all_translated_text += f"[翻译失败] {original_text} | "
                    all_translated_text += "\n"
                all_translated_text += "\n\n"
            
            # 保存文本版本
            with open(txt_output_path, 'w', encoding='utf-8') as f:
                f.write(all_translated_text)
            
            # print(f"DOCX文件翻译完成，原始DOCX文件已保存到: {output_path}")
            print(f"翻译文本已保存到: {txt_output_path}")
            
            return all_translated_text, output_path, txt_output_path
            
        except ImportError:
            # 如果未安装python-docx，则无法处理DOCX
            print("错误: 未安装python-docx库，无法处理DOCX文件。")
            print("请安装: pip install python-docx")
            
            # 尝试生成文本版本
            try:
                with open(input_path, 'r', encoding='utf-8', errors='ignore') as f:
                    text = f.read()
                    translated_text = translate_long_text(text, translator)
                    with open(txt_output_path, 'w', encoding='utf-8') as out:
                        out.write(translated_text)
                    return translated_text, output_path, txt_output_path
            except Exception:
                pass
            
            raise ImportError("需要python-docx库来处理DOCX文件。请安装: pip install python-docx")
        
    except Exception as e:
        raise Exception(f"DOCX文件翻译失败: {str(e)}")

def translate_epub_file(input_path, output_path=None, source_lang='auto', target_lang='zh-CN', progress_callback=None):
    """
    翻译EPUB电子书文件（翻译后仍然为EPUB格式）
    
    参数:
        input_path (str): 输入EPUB文件路径
        output_path (str, optional): 输出文件路径，默认为None（自动生成）
        source_lang (str): 源语言代码，默认为auto自动检测
        target_lang (str): 目标语言代码，默认为zh-CN
        progress_callback (callable): 进度回调函数，接收进度百分比作为参数
        
    返回:
        tuple: (翻译后的文本, 输出文件路径, 文本文件路径)
    """
    try:
        # 更新进度：0%
        if progress_callback:
            progress_callback(0)
            
        # 确保输入文件存在
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"输入文件不存在: {input_path}")
        
        # 如果未指定输出路径，则自动生成
        if output_path is None:
            input_file = Path(input_path)
            filename = input_file.stem
            output_path = os.path.join("output_files", f"{filename}_{target_lang}.epub")
        elif not output_path.lower().endswith('.epub'):
            output_path += '.epub'
        
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 准备文本输出路径
        txt_output_path = output_path.replace('.epub', '.txt')
        all_translated_text = ""
        
        # 临时目录用于解包EPUB
        temp_dir = Path(f"temp_epub_{Path(input_path).stem}")
        output_temp_dir = Path(f"temp_epub_output_{Path(input_path).stem}")
        
        try:
            # 创建临时目录
            temp_dir.mkdir(exist_ok=True)
            output_temp_dir.mkdir(exist_ok=True)
            
            # 解包EPUB
            with zipfile.ZipFile(input_path, 'r') as zip_ref:
                zip_ref.extractall(temp_dir)
            
            # 复制原始文件到输出临时目录（保留所有元数据和结构）
            shutil.copytree(temp_dir, output_temp_dir, dirs_exist_ok=True)
            
            # 更新进度：10%
            if progress_callback:
                progress_callback(10)
            
            # 查找HTML/XHTML文件
            html_files = []
            for file_path in temp_dir.glob('**/*'):
                if file_path.is_file() and file_path.suffix.lower() in ['.html', '.xhtml', '.htm', '.xml']:
                    html_files.append(file_path)
            
            # 如果没有HTML文件，可能不是有效的EPUB
            if not html_files:
                raise ValueError("EPUB文件中未找到HTML内容")
            
            # 创建翻译器
            translator = GoogleTranslator(source=source_lang, target=target_lang)
            
            # 跟踪成功翻译的文件数
            successful_files = 0
            total_files = len(html_files)
            
            # 处理每个HTML文件
            for i, html_file in enumerate(html_files):
                try:
                    # 计算并更新进度（从10%到90%）
                    if progress_callback:
                        progress = 10 + int(80 * (i / total_files))
                        progress_callback(progress)
                    
                    # 获取相对路径，用于在输出临时目录中找到对应文件
                    rel_path = html_file.relative_to(temp_dir)
                    output_html_file = output_temp_dir / rel_path
                    
                    # 尝试确定文件编码
                    encodings = ['utf-8', 'latin-1', 'cp1252', 'iso-8859-1']
                    html_content = None
                    used_encoding = 'utf-8'  # 默认编码
                    
                    for encoding in encodings:
                        try:
                            with open(html_file, 'r', encoding=encoding) as f:
                                html_content = f.read()
                                used_encoding = encoding
                                break
                        except UnicodeDecodeError:
                            continue
                    
                    if html_content is None:
                        print(f"警告: 无法解码文件 {html_file}，跳过处理")
                        all_translated_text += f"[无法解码] 文件: {rel_path}\n\n"
                        continue
                    
                    # 根据文件扩展名选择正确的解析器
                    if html_file.suffix.lower() in ['.xml', '.xhtml']:
                        parser = 'xml'
                    else:
                        parser = 'html.parser'
                    
                    # 解析HTML
                    soup = BeautifulSoup(html_content, parser)
                    
                    # 记录原始文件内容
                    file_text = []
                    
                    # 添加文件标识
                    all_translated_text += f"\n--- 文件: {rel_path} ---\n\n"
                    file_modified = False  # 标记是否对文件进行了修改
                    
                    # 定义一个辅助函数来检查是否应该翻译元素
                    def should_translate_element(element):
                        # 跳过导航、元数据、目录等特殊内容
                        if element.parent and element.parent.name in ['nav', 'metadata', 'guide', 'toc']:
                            return False
                        # 跳过空元素
                        if not element.get_text().strip():
                            return False
                        # 跳过太短的内容（可能是标签、按钮文本等）
                        if len(element.get_text()) < 10:
                            return False
                        # 跳过只包含数字、符号的内容
                        text = element.get_text().strip()
                        if re.match(r'^[\d\s\W]+$', text):
                            return False
                        return True
                    
                    # 找到并翻译所有文本节点（仅翻译段落、标题和列表项中的文本）
                    for element in soup.find_all(['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'li']):
                        if not should_translate_element(element):
                            continue
                            
                        text = element.get_text().strip()
                        if text:  # 只处理非空文本
                            # 记录原始文本
                            file_text.append(text)
                            
                            try:
                                # 添加延迟以避免触发Google翻译的限制
                                if len(file_text) % 5 == 0:
                                    import time
                                    time.sleep(1.5)  # 每处理5个文本块添加1.5秒延迟
                                
                                # 翻译文本
                                translated = translate_long_text(text, translator)
                                
                                # 在文本文件中记录翻译结果
                                all_translated_text += f"原文: {text[:100]}...\n翻译: {translated}\n\n"
                                
                                # 更新HTML元素内容 - 替换元素中的文本，但保留其他内容
                                # 这种方式可以保留HTML标签和结构
                                element.string = translated
                                file_modified = True
                                
                            except Exception as e:
                                print(f"警告: 无法翻译文本: {str(e)[:100]}...")
                                all_translated_text += f"[翻译失败] {text[:200]}...\n\n"
                    
                    # 如果文件被修改，保存修改后的文件
                    if file_modified:
                        with open(output_html_file, 'w', encoding=used_encoding) as f:
                            # 使用prettyprinter保持格式一致
                            f.write(str(soup))
                        successful_files += 1
                    else:
                        # 如果没有修改，直接复制原文件
                        shutil.copy2(html_file, output_html_file)
                    
                except Exception as e:
                    print(f"警告: 处理HTML文件失败 {html_file}: {str(e)}")
                    all_translated_text += f"[处理失败] 文件: {rel_path}\n错误: {str(e)}\n\n"
                    # 确保即使处理失败，原始文件仍被复制
                    if html_file.exists() and not output_html_file.exists():
                        shutil.copy2(html_file, output_html_file)
            
            # 添加摘要信息
            summary = (
                f"EPUB文件翻译摘要:\n"
                f"总文件数: {len(html_files)}\n"
                f"成功翻译: {successful_files}\n"
                f"未能翻译: {len(html_files) - successful_files}\n\n"
            )
            all_translated_text = summary + all_translated_text
            
            # 保存文本版本
            with open(txt_output_path, 'w', encoding='utf-8') as f:
                f.write(all_translated_text)
            
            # 打包为新的EPUB文件
            with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for file_path in output_temp_dir.glob('**/*'):
                    if file_path.is_file():
                        arcname = file_path.relative_to(output_temp_dir)
                        zipf.write(file_path, arcname)
            
            # 更新进度：100%
            if progress_callback:
                progress_callback(100)
                
            print(f"翻译的EPUB文件已保存到: {output_path}")
            print(f"翻译摘要已保存到: {txt_output_path}")
            
            # 返回结果
            return all_translated_text, output_path, txt_output_path
            
        finally:
            # 清理临时目录
            if temp_dir.exists():
                shutil.rmtree(temp_dir)
            if output_temp_dir.exists():
                shutil.rmtree(output_temp_dir)
        
    except Exception as e:
        # 更新进度：失败
        if progress_callback:
            progress_callback(-1)  # 使用-1表示失败
        raise Exception(f"EPUB文件翻译失败: {str(e)}") 